#!/usr/bin/env python3
"""
超导腔性能监控脚本
基于 pyepics 的可配置监控系统，用于测量自动加载时长和闭环稳定时长
"""

import epics
import configparser
import csv
import threading
import time
from datetime import datetime
from enum import Enum
from dataclasses import dataclass
from typing import Dict, List, Optional
import os
import sys


class CavityState(Enum):
    """腔体状态枚举"""
    IDLE = "idle"
    LOADING = "loading"
    CLOSED_LOOP = "closed_loop"


@dataclass
class EventRecord:
    """事件记录数据类"""
    device_name: str
    start_time: datetime
    end_time: datetime
    duration_seconds: float


class CavityMonitor:
    """单个腔体监控器"""
    
    def __init__(self, device_name: str, pv_prefix: str, logger):
        self.device_name = device_name
        self.pv_prefix = pv_prefix
        self.logger = logger
        self.state = CavityState.IDLE
        self.load_start_time = None
        self.uptime_start_time = None
        
        # 创建PV连接
        self.apl_start_pv = epics.PV(f"{pv_prefix}APL_Start")
        self.rf_ready_pv = epics.PV(f"{pv_prefix}RFReady")
        
        # 设置回调函数
        self.apl_start_pv.add_callback(self._on_apl_start_change)
        self.rf_ready_pv.add_callback(self._on_rf_ready_change)
        
        self.logger.log_info(f"[{self.device_name}] 监控器初始化完成")
    
    def _on_apl_start_change(self, pvname=None, value=None, timestamp=None, **kwargs):
        """APL_Start PV变化回调"""
        if value == 1 and self.rf_ready_pv.value == 0:
            self._start_loading_timer(timestamp)
    
    def _on_rf_ready_change(self, pvname=None, value=None, timestamp=None, **kwargs):
        """RFReady PV变化回调"""
        if value == 1:
            self._on_rf_ready_high(timestamp)
        elif value == 0:
            self._on_rf_ready_low(timestamp)
    
    def _start_loading_timer(self, timestamp):
        """开始加载计时"""
        if self.state == CavityState.IDLE:
            self.state = CavityState.LOADING
            self.load_start_time = datetime.fromtimestamp(timestamp)
            self.logger.log_info(f"[{self.device_name}] 检测到加载指令，开始计时...")
    
    def _on_rf_ready_high(self, timestamp):
        """RFReady变为1的处理"""
        end_time = datetime.fromtimestamp(timestamp)
        
        # 如果正在加载，记录加载完成
        if self.state == CavityState.LOADING and self.load_start_time:
            duration = (end_time - self.load_start_time).total_seconds()
            record = EventRecord(
                device_name=self.device_name,
                start_time=self.load_start_time,
                end_time=end_time,
                duration_seconds=duration
            )
            self.logger.log_auto_load(record)
            self.logger.log_success(f"[{self.device_name}] 加载成功, 耗时 {duration:.2f} 秒.")
        
        # 开始闭环计时
        self.state = CavityState.CLOSED_LOOP
        self.uptime_start_time = end_time
        self.logger.log_info(f"[{self.device_name}] 进入闭环状态，开始稳定计时.")
    
    def _on_rf_ready_low(self, timestamp):
        """RFReady变为0的处理"""
        if self.state == CavityState.CLOSED_LOOP and self.uptime_start_time:
            end_time = datetime.fromtimestamp(timestamp)
            duration = (end_time - self.uptime_start_time).total_seconds()
            record = EventRecord(
                device_name=self.device_name,
                start_time=self.uptime_start_time,
                end_time=end_time,
                duration_seconds=duration
            )
            self.logger.log_uptime(record)
            self.logger.log_warning(f"[{self.device_name}] 发生失锁! 闭环稳定时长: {duration:.2f} 秒.")
        
        # 重置状态
        self.state = CavityState.IDLE
        self.load_start_time = None
        self.uptime_start_time = None


class DataLogger:
    """数据记录器"""
    
    def __init__(self, auto_load_file: str, uptime_file: str):
        self.auto_load_file = auto_load_file
        self.uptime_file = uptime_file
        self.lock = threading.Lock()
        
        # 初始化CSV文件
        self._init_csv_files()
    
    def _init_csv_files(self):
        """初始化CSV文件，写入表头"""
        headers = ['DeviceName', 'StartTime', 'EndTime', 'Duration_seconds']
        
        for filename in [self.auto_load_file, self.uptime_file]:
            if not os.path.exists(filename):
                with open(filename, 'w', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    writer.writerow(headers)
    
    def log_auto_load(self, record: EventRecord):
        """记录自动加载时长"""
        self._write_record(self.auto_load_file, record)
    
    def log_uptime(self, record: EventRecord):
        """记录闭环稳定时长"""
        self._write_record(self.uptime_file, record)
    
    def _write_record(self, filename: str, record: EventRecord):
        """写入记录到CSV文件"""
        with self.lock:
            with open(filename, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow([
                    record.device_name,
                    record.start_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3],
                    record.end_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3],
                    f"{record.duration_seconds:.2f}"
                ])
    
    def log_info(self, message: str):
        """记录信息日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{timestamp}] [INFO] {message}")
    
    def log_success(self, message: str):
        """记录成功日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{timestamp}] [SUCCESS] {message}")
    
    def log_warning(self, message: str):
        """记录警告日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"[{timestamp}] [WARNING] {message}")


class SRFCavityMonitorSystem:
    """超导腔监控系统主类"""
    
    def __init__(self, config_file: str = 'config.ini'):
        self.config_file = config_file
        self.config = configparser.ConfigParser()
        self.monitors: List[CavityMonitor] = []
        self.logger = None
        
        # 读取配置
        self._load_config()
        
        # 初始化日志记录器
        auto_load_file = self.config.get('global', 'auto_load_log_file')
        uptime_file = self.config.get('global', 'uptime_log_file')
        self.logger = DataLogger(auto_load_file, uptime_file)
        
        # 创建监控器
        self._create_monitors()
    
    def _load_config(self):
        """加载配置文件"""
        if not os.path.exists(self.config_file):
            raise FileNotFoundError(f"配置文件 {self.config_file} 不存在")
        
        self.config.read(self.config_file, encoding='utf-8')
        print(f"配置文件 {self.config_file} 加载成功")
    
    def _create_monitors(self):
        """根据配置创建监控器"""
        for section_name in self.config.sections():
            if section_name == 'global':
                continue
            
            section = self.config[section_name]
            
            # 检查是否启用
            if section.getboolean('enabled', fallback=False):
                cavity_count = section.getint('cavity_count')
                pv_pattern = section.get('pv_prefix_pattern')
                
                # 为每个腔创建监控器
                for cav_num in range(1, cavity_count + 1):
                    device_name = f"{section_name}-LLRF{cav_num:02d}"
                    pv_prefix = pv_pattern.format(cav_num=cav_num)
                    
                    monitor = CavityMonitor(device_name, pv_prefix, self.logger)
                    self.monitors.append(monitor)
        
        self.logger.log_info(f"共创建 {len(self.monitors)} 个腔体监控器")
    
    def start_monitoring(self):
        """开始监控"""
        self.logger.log_info("超导腔性能监控系统启动")
        self.logger.log_info(f"监控设备数量: {len(self.monitors)}")
        
        try:
            while True:
                time.sleep(1)  # 主循环保持运行
        except KeyboardInterrupt:
            self.logger.log_info("收到停止信号，正在关闭监控系统...")
            self.stop_monitoring()
    
    def stop_monitoring(self):
        """停止监控"""
        self.logger.log_info("超导腔性能监控系统已停止")


def main():
    """主函数"""
    config_file = 'config.ini'
    
    # 检查配置文件
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
    
    try:
        # 创建监控系统
        monitor_system = SRFCavityMonitorSystem(config_file)
        
        # 开始监控
        monitor_system.start_monitoring()
        
    except Exception as e:
        print(f"错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()